package TongJiChengXu

import (
	"regexp"
	"sort"
	"strings"
	"sync"
)

// 分词函数，处理英文大小写和标点
func tokenize(text string) []string {
	// 转换为小写
	text = strings.ToLower(text)
	// 去除标点符号
	reg := regexp.MustCompile(`[^\w\s]`)
	text = reg.ReplaceAllString(text, "")
	// 按空格分词
	return strings.Fields(text)
}

// WordFrequency 基础词频统计函数
func WordFrequency(text string) map[string]int {
	words := tokenize(text)
	freq := make(map[string]int)
	for _, word := range words {
		freq[word]++
	}
	return freq
}

// WordCount 单词频率结构体
type WordCount struct {
	Word  string
	Count int
}

// sortWordFrequency 按频率排序的函数
func sortWordFrequency(freq map[string]int) []WordCount {
	wordCounts := make([]WordCount, 0, len(freq))
	for word, count := range freq {
		wordCounts = append(wordCounts, WordCount{Word: word, Count: count})
	}
	sort.Slice(wordCounts, func(i, j int) bool {
		return wordCounts[i].Count > wordCounts[j].Count
	})
	return wordCounts
}

// WordFrequencyConcurrent 并发分词并统计词频
func WordFrequencyConcurrent(text string) map[string]int {
	lines := strings.Split(text, "\n")
	freq := make(map[string]int)
	var mu sync.Mutex
	var wg sync.WaitGroup

	for _, line := range lines {
		wg.Add(1)
		go func(line string) {
			defer wg.Done()
			words := tokenize(line)
			localFreq := make(map[string]int)
			for _, word := range words {
				localFreq[word]++
			}
			mu.Lock()
			for word, count := range localFreq {
				freq[word] += count
			}
			mu.Unlock()
		}(line)
	}
	wg.Wait()
	return freq
}

// TrieNode 定义
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
	count    int
}

// Trie 树结构
type Trie struct {
	root *TrieNode
}

// NewTrie 初始化Trie
func NewTrie() *Trie {
	return &Trie{root: &TrieNode{children: make(map[rune]*TrieNode)}}
}

// Insert 插入单词
func (t *Trie) Insert(word string) {
	node := t.root
	for _, char := range word {
		if node.children[char] == nil {
			node.children[char] = &TrieNode{children: make(map[rune]*TrieNode)}
		}
		node = node.children[char]
	}
	node.isEnd = true
	node.count++
}

// Search 查找单词频率
func (t *Trie) Search(word string) int {
	node := t.root
	for _, char := range word {
		if node.children[char] == nil {
			return 0
		}
		node = node.children[char]
	}
	if node.isEnd {
		return node.count
	}
	return 0
}

// WordFrequencyTrie 使用Trie树统计词频
func WordFrequencyTrie(text string) *Trie {
	trie := NewTrie()
	words := tokenize(text)
	for _, word := range words {
		trie.Insert(word)
	}
	return trie
}
